import torch
import matplotlib.pyplot as plt

# 示例张量
load_tensors_1 = torch.cat(
    [torch.load(f"../ffn/second_line_input/tensor{i}.pth", weights_only=True) for i in range(32)])
load_tensors_2 = torch.matmul(load_tensors_1, torch.load("../1-1-2_target_657.pth", weights_only=True).to("cpu"))


def show_plot(data, title):
    # 对第三维度进行降采样和平均
    stride = 100
    data_downsampled = data.unfold(dimension=2, size=stride, step=stride).mean(dim=-1)  # [32, 7, 1100]

    # 创建子图网格布局
    fig, axs = plt.subplots(8, 4, figsize=(20, 16))  # 8 行 4 列，适合 32 个子图
    fig.suptitle(title, fontsize=40)
    # 遍历每个样本并绘制到子图
    for i, ax in enumerate(axs.flat):
        sample_data = data_downsampled[i].detach().numpy()  # 当前样本数据，形状 [7, 1100]

        # 绘制热力图
        im = ax.imshow(sample_data, cmap='viridis', aspect='auto')

        # 设置子图标题
        ax.set_title(f"Block {i + 1}", fontsize=18)

        # 隐藏轴标签
        ax.set_xticks([])
        ax.set_yticks([])

    # 添加一个全局颜色条
    fig.colorbar(im, ax=axs, orientation='horizontal', fraction=0.02, pad=0.04, label="Value")

    # 调整布局
    # plt.tight_layout()
    save_path = f"{title}-热力图"
    if save_path:
        plt.savefig(save_path, dpi=300)  # 保存为高分辨率图片
        print(f"Figure saved to {save_path}")

    # 显示图形
    plt.show()



if __name__ == '__main__':
    show_plot(load_tensors_1, "feature map旋转前")
    show_plot(load_tensors_2, "feature map旋转后")